package problems.maze20;

import problems.RewardFunction;

public class Maze20Rewards implements RewardFunction<Integer,Integer>{

	private boolean threeObjectives;
	
	public Maze20Rewards(boolean threeObj){
		this.threeObjectives = threeObj;
	}

	@Override
	public int numberOfObjectives() {
		if(threeObjectives){
			return 3;
		} else {
			return 2;
		}
	}

	@Override
	public double[] expectedReward(Integer state, Integer action,
			Integer statePrime) {
		
		if(this.threeObjectives){ 
			if(state.intValue() == 7){//TERMINAL
				double[] rew = {0,0,0};
				return rew;
			} else if (statePrime.intValue() == 7){//GOAL REACHED
				double[] rew = {0,0,150};//150 following Hauskrecht 2000
				return rew;
			}
			
			if(action.intValue() == Maze20Actions.SENSEEW || 
					action.intValue() == Maze20Actions.SENSENS ){
				double[] rew = {-1, 2, 0};
				return rew;
			} else {
				if(state.intValue() == statePrime.intValue()){
					//It bumped into the wall
					double[] rew = {0, 0, 0};
					return rew;
				} else { //Valid move
					double[] rew = {0, 4, 0};
					return rew;
				}
			}
		} else { // two objectives
			if(state.intValue() == 7){//TERMINAL
				double[] rew = {0,0};
				return rew;
			} else if (statePrime.intValue() == 7){//GOAL REACHED
				double[] rew = {0,150};//150 following Hauskrecht 2000
				return rew;
			}
			
			if(action.intValue() == Maze20Actions.SENSEEW || 
					action.intValue() == Maze20Actions.SENSENS ){
				double[] rew = {2, 0};
				return rew;
			} else {
				if(state.intValue() == statePrime.intValue()){
					//It bumped into the wall
					double[] rew = {0, 0};
					return rew;
				} else { //just a valid move
					double[] rew = {4, 0};
					return rew;
				}
			}
		}
	}

	@Override
	public double[] expectedReward(Integer state, Integer action) {
		if(this.threeObjectives){ 
			if(state.intValue() == 7){//TERMINAL
				double[] rew = {0,0,0};
				return rew;
			} else if (state.intValue() == 12 && 
					action.intValue()==Maze20Actions.SOUTH){//GOAL REACHABLE
				double[] rew = {0,0,150};//following Hauskrecht 2000
				rew[2] *= 0.7;
				rew[1] = -0.3;
				return rew;
			}
			if(action.intValue() == Maze20Actions.SENSEEW || 
					action.intValue() == Maze20Actions.SENSENS ){
				double[] rew = {-1, 2, 0};
				return rew;
			} else {
				double bumpProb = 
						Maze20Transitions.getTransitions().probabilityOfTransition(state, action, state);
				//It bumped into the wall
				double[] rew = {0, 4, 0};
				rew[1] *= (1-bumpProb);
				return rew;
			}
		} else { // two objectives
			if(state.intValue() == 7){//TERMINAL
				double[] rew = {0,0};
				return rew;
			} else if (state.intValue() == 12 && 
					action.intValue()==Maze20Actions.SOUTH){//GOAL REACHABLE
				double[] rew = {0,150};
				rew[1] *= 0.7;
				rew[0] = 0;
				return rew;
			}
			if(action.intValue() == Maze20Actions.SENSEEW || 
					action.intValue() == Maze20Actions.SENSENS ){
				double[] rew = {2, 0};//folloing Hauskrecht 
				return rew;
			} else {
				double bumpProb = 
						Maze20Transitions.getTransitions().probabilityOfTransition(state, action, state);
				//It bumped into the wall
				double[] rew = {4, 0};//Following Hauskrecht
				rew[0] *= (1-bumpProb);
				return rew;
			}
		}
	}

	@Override
	public double[] getReward(Integer state, Integer action, Integer statePrime) {
		//Them deterministic
		return this.expectedReward(state, action, statePrime);
	}

	@Override
	public double[] getReward(Integer state, Integer action) {
		//Them deterministic
		return this.expectedReward(state, action);
	}

	@Override
	public double[] minRewards() {
		if(this.threeObjectives){
			double[] rew = {-1, 0, 0};
			return rew;
		} else {
			double[] rew = {0, 0};
			return rew;
		}
	}
	
	
}
